package com.study.mirrorforest.main.common.utils.qwen;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.study.mirrorforest.main.config.QwenConfig;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.concurrent.TimeUnit;

/**
 * Qwen API 调用工具类
 * 封装了与阿里云百炼 Qwen API 的交互逻辑
 * @author AI Assistant
 */
@Component
public class QwenApiUtil {

    @Autowired
    private QwenConfig qwenConfig;

    private final OkHttpClient httpClient;
    private final ObjectMapper objectMapper;

    public QwenApiUtil() {
        // 初始化 OkHttpClient
        this.httpClient = new OkHttpClient.Builder()
                .connectTimeout(30, TimeUnit.SECONDS)
                .readTimeout(60, TimeUnit.SECONDS)
                .writeTimeout(30, TimeUnit.SECONDS)
                .build();
        this.objectMapper = new ObjectMapper();
    }

    /**
     * 调用 Qwen API 生成文本
     * @param requestDTO 请求参数
     * @return Qwen API 响应
     * @throws QwenApiException 调用失败时抛出异常
     */
    public QwenResponseDTO generateText(QwenRequestDTO requestDTO) throws QwenApiException {
        // 1. 参数校验
        validateRequest(requestDTO);

        // 2. 构建请求体
        String requestBody = buildRequestBody(requestDTO);

        // 3. 发送 HTTP 请求
        try {
            Request request = new Request.Builder()
                    .url(qwenConfig.getEndpoint())
                    .addHeader("Authorization", "Bearer " + qwenConfig.getKey())
                    .addHeader("Content-Type", "application/json")
                    .post(RequestBody.create(requestBody, MediaType.parse("application/json")))
                    .build();

            try (Response response = httpClient.newCall(request).execute()) {
                if (!response.isSuccessful()) {
                    String errorBody = response.body() != null ? response.body().string() : "Unknown error";
                    throw new QwenApiException("Qwen API 调用失败，HTTP 状态码: " + response.code() + ", 错误信息: " + errorBody);
                }

                // 4. 解析响应
                String responseBody = response.body().string();
                return parseResponse(responseBody);
            }
        } catch (IOException e) {
            throw new QwenApiException("网络请求失败: " + e.getMessage(), e);
        }
    }

    /**
     * 快捷方法：只传入 prompt 文本
     * @param prompt 用户问题
     * @return Qwen API 响应
     * @throws QwenApiException 调用失败时抛出异常
     */
    public QwenResponseDTO chat(String prompt) throws QwenApiException {
        return generateText(new QwenRequestDTO(prompt));
    }

    /**
     * 快捷方法：传入 prompt 和系统提示词
     * @param prompt 用户问题
     * @param systemPrompt 系统提示词
     * @return Qwen API 响应
     * @throws QwenApiException 调用失败时抛出异常
     */
    public QwenResponseDTO chat(String prompt, String systemPrompt) throws QwenApiException {
        QwenRequestDTO request = new QwenRequestDTO(prompt);
        request.setSystemPrompt(systemPrompt);
        return generateText(request);
    }

    /**
     * 校验请求参数
     */
    private void validateRequest(QwenRequestDTO requestDTO) throws QwenApiException {
        if (requestDTO == null) {
            throw new QwenApiException("请求参数不能为空");
        }
        if (requestDTO.getPrompt() == null || requestDTO.getPrompt().trim().isEmpty()) {
            throw new QwenApiException("用户输入不能为空");
        }
        if (qwenConfig.getKey() == null || qwenConfig.getKey().isEmpty()) {
            throw new QwenApiException("Qwen API Key 未配置，请在 application.properties 中配置 qwen.api.key");
        }
    }

    /**
     * 构建请求体 JSON
     */
    private String buildRequestBody(QwenRequestDTO requestDTO) {
        try {
            StringBuilder messagesBuilder = new StringBuilder();
            messagesBuilder.append("[");

            // 添加系统提示词（如果有）
            if (requestDTO.getSystemPrompt() != null && !requestDTO.getSystemPrompt().isEmpty()) {
                messagesBuilder.append("{\"role\":\"system\",\"content\":\"")
                        .append(escapeJson(requestDTO.getSystemPrompt()))
                        .append("\"},");
            }

            // 添加用户输入
            messagesBuilder.append("{\"role\":\"user\",\"content\":\"")
                    .append(escapeJson(requestDTO.getPrompt()))
                    .append("\"}");

            messagesBuilder.append("]");

            // 构建完整的请求体
            return String.format(
                    "{\"model\":\"%s\",\"input\":{\"messages\":%s},\"parameters\":{\"result_format\":\"message\",\"max_tokens\":%d,\"temperature\":%.2f}}",
                    qwenConfig.getModel(),
                    messagesBuilder.toString(),
                    requestDTO.getMaxTokens() != null ? requestDTO.getMaxTokens() : qwenConfig.getMaxTokens(),
                    requestDTO.getTemperature() != null ? requestDTO.getTemperature() : qwenConfig.getTemperature()
            );
        } catch (Exception e) {
            throw new RuntimeException("构建请求体失败: " + e.getMessage(), e);
        }
    }

    /**
     * 解析 API 响应
     */
    private QwenResponseDTO parseResponse(String responseBody) throws QwenApiException {
        try {
            JsonNode root = objectMapper.readTree(responseBody);

            // 检查是否有错误
            if (root.has("code") && !root.get("code").asText().isEmpty()) {
                String errorCode = root.get("code").asText();
                String errorMessage = root.has("message") ? root.get("message").asText() : "Unknown error";
                throw new QwenApiException("Qwen API 返回错误，错误码: " + errorCode + ", 错误信息: " + errorMessage);
            }

            // 提取响应数据
            QwenResponseDTO response = new QwenResponseDTO();

            // 提取生成的文本
            if (root.has("output") && root.get("output").has("choices")) {
                JsonNode choices = root.get("output").get("choices");
                if (choices.isArray() && choices.size() > 0) {
                    JsonNode firstChoice = choices.get(0);
                    if (firstChoice.has("message") && firstChoice.get("message").has("content")) {
                        response.setText(firstChoice.get("message").get("content").asText());
                    }
                    if (firstChoice.has("finish_reason")) {
                        response.setFinishReason(firstChoice.get("finish_reason").asText());
                    }
                }
            }

            // 提取 request_id
            if (root.has("request_id")) {
                response.setRequestId(root.get("request_id").asText());
            }

            // 提取 token 使用量
            if (root.has("usage") && root.get("usage").has("total_tokens")) {
                response.setTotalTokens(root.get("usage").get("total_tokens").asInt());
            }

            return response;
        } catch (IOException e) {
            throw new QwenApiException("解析 API 响应失败: " + e.getMessage(), e);
        }
    }

    /**
     * 转义 JSON 字符串中的特殊字符
     */
    private String escapeJson(String str) {
        return str.replace("\\", "\\\\")
                .replace("\"", "\\\"")
                .replace("\n", "\\n")
                .replace("\r", "\\r")
                .replace("\t", "\\t");
    }
}

